import matplotlib.pyplot as plt
from sklearn.tree import DecisionTreeClassifier
from sklearn.ensemble import RandomForestClassifier
from sklearn.model_selection import train_test_split
import pandas as pd
import numpy as np
from sklearn import tree

df = pd.read_excel('data2.xlsx')
Ycol = ['y']
Xcols = list(set(df.columns) - set(Ycol))
needToNumCols = {'marital':['single', 'divorced', 'married'], 'housing':['yes', 'no'], 'pre_outcome':['success', 'nonexistent', 'failure'],
                 'contact':['telephone', 'cellular'], 'loan':['yes', 'no', 'unknown']}

def strColToNum(strCol, strMapping):
    sub = 0
    for i in strMapping:
        strCol = strCol.replace(i, sub)
        sub += 1
    return strCol

def doStrColToNum(df):
    for colName, strMapping in needToNumCols.items():
        df[colName] = strColToNum(df[colName], strMapping)
doStrColToNum(df)

clf = DecisionTreeClassifier()
Xdata = df[Xcols]
Ydata = df['y']
clf.fit(Xdata, Ydata)

df = pd.read_csv('学号2.csv')
doStrColToNum(df)
print('y')
ret = clf.predict(df[Xcols])
for i in ret:
    print(i)